import librosa
import numpy as np
from sklearn.neighbors import KNeighborsClassifier
import os

def extract_mfc_from_file(file_path):
    """
    从给定音频文件中提取MFCC特征
    """
    y, sr = librosa.load(file_path, sr=None)
    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
    mfcc_mean = np.mean(mfcc.T, axis=0)
    return mfcc_mean

# 训练集文件夹路径，里面分别有'你'和'好'的子文件夹，子文件夹中存放对应音频
train_folder_path = r'E:\实训\音频数据集'
# 预测的音频所在文件夹路径
predict_folder_path = r'E:\实训\测试数据集'

# 提取训练集的MFCC特征
mfc_ni = [extract_mfc_from_file(os.path.join(root, file)) for root, dirs, files in os.walk(os.path.join(train_folder_path, '你')) for file in files]
mfc_hao = [extract_mfc_from_file(os.path.join(root, file)) for root, dirs, files in os.walk(os.path.join(train_folder_path, '好')) for file in files]

X = np.array(mfc_ni + mfc_hao)
# 对应标签，假设'你'文件夹中的音频对应标签为'你'，'好'文件夹中的音频对应标签为'好'
y = ['你'] * len(mfc_ni) + ['好'] * len(mfc_hao)

# 构建KNN模型
knn = KNeighborsClassifier(n_neighbors=1)
knn.fit(X, y)

# 提取预测文件夹中音频的MFCC特征并进行预测
for root, dirs, files in os.walk(predict_folder_path):
    for file in files:
        file_path = os.path.join(root, file)
        mfcc_to_predict = extract_mfc_from_file(file_path)
        # 确保mfcc_to_predict是一个二维数组，形状为(1, n_features)
        mfcc_to_predict_2d = mfcc_to_predict.reshape(1, -1)
        prediction = knn.predict(mfcc_to_predict_2d)  # 直接传递二维数组
        print(f'Prediction for {file}: {prediction[0]}')